{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1, Topic 1, Lab A: Resynchronizing Traces with Sum of Absolute Difference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "NOTE: This lab references some (commercial) training material on [ChipWhisperer.io](https://www.ChipWhisperer.io). You can freely execute and use the lab per the open-source license (including using it in your own courses if you distribute similarly), but you must maintain notice about this source location. Consider joining our training course to enjoy the full experience.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY:** *At the end of SCA101, you saw how the communication lines of a microcontroller can be used as a trigger in a side channel attack. This attack was made much more difficult by the presence of jitter; the traces didn't all line up with each other. One thing you might have asked after that lab was if jitter can be used as a countermeasure to a CPA attack. If our target introduces enough jitter, will our CPA attack become impractical?*\n",
    "\n",
    "*In this lab, we'll look at how jitter-based countermeasures can be overcome by resynchronizing the traces. More specifically, we'll use the sum of absolute difference (SAD) measure that we last looked at early on in SCA101.* \n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "* Attempting a CPA attack againt a jittery AES implementation\n",
    "* Writing code to resynchronize traces using SAD\n",
    "* Using ChipWhisperer Analyzer to resynchronize traces."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "This lab will build upon the material in SCA101. In particular, you may want to have a quick look over the following labs:\n",
    "\n",
    "* Lab 2_1 (We used SAD to crack the password here)\n",
    "* Lab 5_1 (You saw how jittery traces can disrupt a CPA attack)\n",
    "* Lab 4_3 (Intro to ChipWhisperer Analyzer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trace Capture\n",
    "\n",
    "We'll be attacking our usual implementation of AES in this lab, with one small modification. If we use the EXTRA_OPTS=ADD_JITTER, a for loop that runs between 0 and 15 times through will be inserted:\n",
    "\n",
    "```C\n",
    "  #ifdef ADD_JITTER\n",
    "  for (volatile uint8_t k = 0; k < (*pt & 0x0F); k++);\n",
    "  #endif\n",
    "```\n",
    "\n",
    "As you can see, this jitter is actually based on the first byte of our plaintext. This in itself is a vulnerability. We won't use this fact, but as an exercise, try hypothocizing some ways to use this to overcome the jitter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Perform the capture, resulting in a project of 200 jittery traces. See the notebooks to copy your data into!\n",
    "#\n",
    "raise NotImplementedError(\"Add your code here, and delete this.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CPA Against Jittery Traces\n",
    "\n",
    "As a first test, let's try plotting a few traces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt = cw.plot([])\n",
    "for i in range(10):\n",
    "    plt *= cw.plot(proj.waves[i])\n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which looks fairly discouraging. Remember that for a CPA attack, we're calculation correlation across our trace set by point, meaning if the SBox output for one trace is in a different spot than the rest, it will decrease the correlation, not increase it. It probably won't end well, but let's try an attack anyway:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(proj, leak_model)\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the attack completely failed. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manual SAD Resync\n",
    "\n",
    "In concept at least, we have a similar situation to the password bypass, except this time the time difference between the traces is an obstacle instead of a vulnerability we're exploiting. We can actually use a very similar technique to figure out how much the traces have shifted! The only modification to the SAD calculation we need to make is that we'll need to use a subset of our reference trace as the reference instead. This means we won't have to figure out how to deal with the parts of the trace that are shifted outside the capture window.\n",
    "\n",
    "To find the offset, all we need to do is slide the reference along the trace we're trying to resyncronize, calculating the SAD at each offset. The difference between the offset we picked for our reference and the offset with the lowest SAD will be our timeshift! In practice, we can skip finding all the offsets and just use the first offset that falls below a threshold like we did in the password bypass. If we captured a much longer power trace that had all the rounds of AES, this would also have the advantage of not accidentally matching with a later round.\n",
    "\n",
    "To start, make a function that figures out the offset between two traces using SAD:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_trace_offset(ref_trace, orig_offset, target_trace, threshold):\n",
    "    ref_len = len(ref_trace)\n",
    "    for i in range(0, len(target_trace) - ref_len):\n",
    "        if SAD_SUM < threshold:\n",
    "            return i - orig_offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(50)\n",
    "b = np.arange(70)[20:]\n",
    "assert calculate_trace_offset(a[35:40], 35, b, 1) == -20\n",
    "assert calculate_trace_offset(a[35:40], 35, b, 29) == -25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's pick a reference trace. First, let's plot the first trace. Some stuff to keep in mind:\n",
    "* We want it to be fairly \"unique\" - we don't want something earlier or later on to also match easily.\n",
    "* The maximum offset, and therefore shift, we can get is limited how close the reference is to the edges of the power trace.\n",
    "* This isn't really a concern for this lab, but we don't want it to be unique per trace. For example, if your trigger was a bit earlier (think the UART lab), the transition between serial communication and AES is very distinct. It's a bad SAD location, however, since the serial part varies a lot between traces\n",
    "\n",
    "for example:\n",
    "\n",
    "![](img/Resync_traces_ref.png)\n",
    "\n",
    "the reference in green would be a good spot to pick, while the reference in red would be a poor choice\n",
    "\n",
    "There's a bit of a learning curve to this, so don't be afraid to come back and select a different reference trace if you find it's not working well when trying to resynchronize other traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(proj.waves[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_trace = proj.waves[0][???:???]\n",
    "cw.plot(proj.waves[0][1700:2000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing we can do is slide our reference along another trace and calculate the SAD at each point. This should be similar to your SAD offset function earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sad_plot(ref_trace, target_trace):\n",
    "    ref_len = len(ref_trace)\n",
    "    sads = []\n",
    "    for i in range(0, len(target_trace) - ref_len):\n",
    "        sads.append(np.sum(abs(ref_trace-target_trace[i:i+ref_len])))\n",
    "    return sads\n",
    "\n",
    "cw.plot(get_sad_plot(ref_trace, proj.waves[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you don't see a single downward spike:\n",
    "\n",
    "![](img/GoodVBadRef.png)\n",
    "\n",
    "instead of a few spikes, you'll want to select a different reference. You can also use this plot to get an idea of what sort of threshold to use.\n",
    "\n",
    "Anyway, try it on some different traces. You should get a varying amount offsets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we can easily get how much our trace is offset by, let's see if the plots actually line up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_offset = calculate_trace_offset(ref_trace, ???, proj.waves[1], 10)\n",
    "print(target_offset)\n",
    "\n",
    "cw.plot(proj.waves[0][:]) * cw.plot(proj.waves[1][target_offset:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks good! Try plotting a few different traces, however, and you'll run into a problem: how do we deal with a trace that needs to be shifted forward? There's a related problem for our CPA attack well: we need all the traces to be the same length, but if we move traces backwards, we won't have any data at the end, meaning the traces will be shorter than usual. The easiest way to deal with these problems is just to fill in the rest of the data with zeros:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "target_offset2 = calculate_trace_offset(ref_trace, ???, proj.waves[5], 10)\n",
    "print(target_offset2)\n",
    "new_trace = np.zeros(len(proj.waves[0]))\n",
    "new_trace[-target_offset2:] = proj.waves[5][:target_offset2]\n",
    "\n",
    "cw.plot(proj.waves[0]) * cw.plot(new_trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Anyway, take what we've done and use it to make a function that takes in a trace and returns one that's resynchronized with the reference.\n",
    "\n",
    "**HINT: You'll need to handle 3 cases: offset > 0, offset < 0, and offset = 0**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def resync_with_SAD(ref, orig_offset, target_trace, threshold):\n",
    "    target_offset = calculate_trace_offset(ref, orig_offset, target_trace, threshold)\n",
    "    if target_offset > 0:\n",
    "        new_trace = np.zeros(len(target_trace))\n",
    "        # uncomment the correct calculation\n",
    "        # new_trace[:-target_offset] = target_trace[target_offset:]\n",
    "        # new_trace[-target_offset:] = target_trace[:target_offset]\n",
    "        return new_trace\n",
    "    elif target_offset < 0:\n",
    "        # uncomment the correct calculation\n",
    "        # new_trace[:-target_offset] = target_trace[target_offset:]\n",
    "        # new_trace[-target_offset:] = target_trace[:target_offset]\n",
    "        return new_trace\n",
    "    return target_trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test it out on some jittery traces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "traces = [resync_with_SAD(ref_trace, ???, proj.waves[i], 10) for i in range(10)]\n",
    "plt = cw.plot([])\n",
    "for i in range(10):\n",
    "    plt *= cw.plot(traces[i])\n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully all your traces lined up perfectly. If not, you might have to go back and select another reference or adjust your `resync_with_SAD()` function.\n",
    "\n",
    "Let's make a new project, resynchronize our traces, and insert them into that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resync_proj = cw.create_project(\"Lab 1_1_Resync\", overwrite=True)\n",
    "for trace in proj.traces:\n",
    "    resync_wave = resync_with_SAD(ref_trace, ???, trace.wave, 10)\n",
    "    new_trace = cw.Trace(resync_wave, trace.textin, trace.textout, trace.key)\n",
    "    resync_proj.traces.append(new_trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and let's retry the attack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(resync_proj, leak_model)\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see the attack succeed this time!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SAD Resync with ChipWhisperer Analyzer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like with the CPA attack itself, ChipWhisperer Analyzer can make our lives a lot easier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resync_traces = cwa.preprocessing.ResyncSAD(proj)\n",
    "resync_traces.ref_trace = 0\n",
    "resync_traces.target_window = (???, ???)\n",
    "resync_traces.max_shift = ???\n",
    "resync_analyzer = resync_traces.preprocess()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All these settings should look familiar except `max_shift`, which will cause the ResyncSAD object to discard a trace if it needs to be shifted more than `max_shift`. Why is this useful? Consider the case where there's enough jitter to move the point we're looking at outside of what we captured. It would be much better to simply discard this trace, rather than including it and having it disrupt our data.\n",
    "\n",
    "Anyway, plotting 10 resychronized traces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt = cw.plot([])\n",
    "for i in range(10):\n",
    "    plt *= cw.plot(resync_analyzer.waves[i])\n",
    "plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should find that everything is resychronized, just as it was from our code. Running the attack again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(resync_analyzer, leak_model)\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which easily succeeds."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions & Next Steps\n",
    "\n",
    "In this lab, you saw how we can resychronize traces using sum of absolute difference, allowing us to overcome jittery traces. With two uses of SAD under your belt, you should start to see how it is a useful metric for comparing two power traces, as it is both simple and effective. If you're looking for another usecase for SAD in the ChipWhisperer project, we also have it available as a trigger for the CW1200. This allows you to trigger based on the shape of a wave. In the case of our simple attack, it would completely remove the need for a trigger pin!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<small>NO-FUN DISCLAIMER: This material is Copyright (C) NewAE Technology Inc., 2015-2020. ChipWhisperer is a trademark of NewAE Technology Inc., claimed in all jurisdictions, and registered in at least the United States of America, European Union, and Peoples Republic of China.\n",
    "\n",
    "Tutorials derived from our open-source work must be released under the associated open-source license, and notice of the source must be *clearly displayed*. Only original copyright holders may license or authorize other distribution - while NewAE Technology Inc. holds the copyright for many tutorials, the github repository includes community contributions which we cannot license under special terms and **must** be maintained as an open-source release. Please contact us for special permissions (where possible).\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.</small>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
